# 导入异步编程模块
import asyncio

# 从 mcp SDK 导入客户端会话、stdio 服务器参数、类型定义
from mcp import ClientSession, StdioServerParameters, types

# 导入 stdio 客户端适配器
from mcp.client.stdio import stdio_client

# 导入 FastMCP 服务器实现
from mcp.server.fastmcp import FastMCP

# 创建一个名为 "Hello-MCP" 的 FastMCP 服务器实例
mcp = FastMCP(name="Hello-MCP")


# 注册一个名为 hello 的工具，返回固定问候语
@mcp.tool()
def hello(name: str) -> str:
    """这是一个hello工具"""
    # 返回一段问候语
    return "Hello from MCP server!"


@mcp.tool()
def query_db(query: str) -> str:
    """查询数据库"""
    # 返回一段问候语
    return "Hello from MCP server!"


# 注册一个资源，URI 为 mcp://info，返回服务器基本信息
@mcp.resource("mcp://info")
async def get_info():
    # 返回服务器的名称、版本和描述信息
    return {
        "name": "Hello-MCP",
        "version": "1.0.0",
        "description": "这是一个MCP 服务器",
    }


# 注册一个 prompt，名称为 welcome，返回欢迎信息
@mcp.prompt("welcome")
async def get_welcome():
    # 返回欢迎信息
    return "欢迎使用 MCP 服务器！我是你的 AI 助手。"


# 定义客户端逻辑，自动启动服务器并进行连通性测试
async def run_client():
    # 构造 stdio 服务器参数，指定用 python 运行本文件并传递 "serve" 参数
    server_params = StdioServerParameters(command="python", args=[__file__, "serve"])

    # 以异步方式连接到 stdio 服务器，获取读写流
    async with stdio_client(server_params) as (read, write):
        # 基于读写流创建客户端会话
        async with ClientSession(read, write) as session:
            # 初始化协议
            await session.initialize()

            # 获取并打印所有工具名称
            tools = await session.list_tools()
            print("[Tools]", [t.name for t in tools.tools])

            # 获取并打印所有资源 URI
            resources = await session.list_resources()
            print("[Resources]", [r.uri for r in resources.resources])

            # 获取并打印所有 prompt 名称
            prompts = await session.list_prompts()
            print("[Prompts]", [p.name for p in prompts.prompts])

            # 调用 hello 工具
            result = await session.call_tool("hello", arguments={})
            # 提取并打印返回的文本内容
            tool_texts = []
            for block in result.content:
                # 判断内容块是否为文本类型
                if isinstance(block, types.TextContent):
                    tool_texts.append(block.text)
            # 打印 hello 工具的返回内容
            print(
                "[Call hello]", " | ".join(tool_texts) or str(result.structuredContent)
            )

            # 获取 resource 内容
            resource_texts = []
            for resource in resources.resources:
                # 读取每个资源的内容
                resource_result = await session.read_resource(resource.uri)
                for block in resource_result.contents:
                    # 判断内容块是否为文本资源类型
                    if isinstance(block, types.TextResourceContents):
                        resource_texts.append(block.text)
                # 打印每个资源的内容
                print(f"[Resource: {resource.uri}]", " | ".join(resource_texts))

            # 获取 prompt 内容
            prompt_texts = []
            for prompt in prompts.prompts:
                # 获取每个 prompt 的内容
                prompt_result = await session.get_prompt(prompt.name)
                for message in prompt_result.messages:
                    # 判断消息是否为 PromptMessage 类型
                    if isinstance(message, types.PromptMessage):
                        prompt_texts.append(message.content.text)
                # 打印每个 prompt 的内容
                print(f"[Prompt: {prompt.name}]", " | ".join(prompt_texts))


# 服务器入口：当以 "serve" 参数运行时，启动服务器
def run_server_stdio():
    # 打印服务器启动信息
    print("启动 MCP 服务器（stdio 模式）...")
    # 以 stdio 方式运行 MCP 服务器
    mcp.run(transport="stdio")


# 主函数：根据参数决定运行模式
def main():
    # 导入 sys 模块以获取命令行参数
    import sys

    # 判断命令行参数是否为 "serve"
    if len(sys.argv) >= 2 and sys.argv[1] == "serve":
        # 服务器模式：运行 stdio 传输
        run_server_stdio()
    else:
        # 客户端模式：启动服务器并连接，然后做连通性检查
        print("启动 MCP 客户端测试...")
        print("将自动启动服务器进程并测试连接")
        # 运行客户端测试逻辑
        asyncio.run(run_client())


# 如果作为主程序运行，则调用 main 函数
if __name__ == "__main__":
    main()
